{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.1307,), (0.3081,))\n",
    "])\n",
    "\n",
    "train_dataset = datasets.MNIST(root=\"./data\", train=True, download=True, transform=transform)\n",
    "test_dataset = datasets.MNIST(root=\"./data\", train=False, download=True, transform=transform)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Model as MyModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Teacher Model...\n",
      "Epoch 1/5, Loss: 0.3139\n",
      "Epoch 2/5, Loss: 0.2172\n",
      "Epoch 3/5, Loss: 0.1790\n",
      "Epoch 4/5, Loss: 0.1662\n",
      "Epoch 5/5, Loss: 0.1640\n"
     ]
    }
   ],
   "source": [
    "def train_teacher(model, train_loader, epochs = 5, lr = 0.01):\n",
    "    model.to(device)\n",
    "    optimizer = optim.Adam(model.parameters(), lr = lr)\n",
    "    loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        total_loss = 0\n",
    "        for images, labels in train_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            logits = model(images)\n",
    "            loss = loss_fn(logits, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_loss += loss.item()\n",
    "        print(f\"Epoch {epoch+1}/{epochs}, Loss: {total_loss / len(train_loader):.4f}\")\n",
    "\n",
    "\n",
    "print(\"Training Teacher Model...\")\n",
    "teacher_model = MyModel.TeacherModel()\n",
    "train_teacher(teacher_model, train_loader, epochs=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Student Model...\n",
      "Epoch 1/5, Loss: 0.6989\n",
      "Epoch 2/5, Loss: 0.4112\n",
      "Epoch 3/5, Loss: 0.3773\n",
      "Epoch 4/5, Loss: 0.3565\n",
      "Epoch 5/5, Loss: 0.3342\n"
     ]
    }
   ],
   "source": [
    "def train_student(student_model, teacher_model, train_loader, epochs=5, lr=0.01, T=3.0, alpha=0.5):\n",
    "    student_model.to(device)\n",
    "    teacher_model.to(device)\n",
    "    teacher_model.eval()\n",
    "    optimizer = optim.Adam(student_model.parameters(), lr=lr)\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        student_model.train()\n",
    "        total_loss = 0\n",
    "        for images, labels in train_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            student_logits = student_model(images)\n",
    "            with torch.no_grad():\n",
    "                teacher_logits = teacher_model(images)\n",
    "\n",
    "            loss = MyModel.distillation_loss(student_logits, teacher_logits, labels, T=T, alpha=alpha)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_loss += loss.item()\n",
    "\n",
    "        print(f\"Epoch {epoch+1}/{epochs}, Loss: {total_loss / len(train_loader):.4f}\")\n",
    "\n",
    "\n",
    "print(\"Training Student Model...\")\n",
    "student_model = MyModel.StudentModel()\n",
    "train_student(student_model, teacher_model, train_loader, epochs=5, T=3.0, alpha=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Teacher Model Accuracy: 0.9532\n",
      "Student Model Accuracy: 0.9585\n"
     ]
    }
   ],
   "source": [
    "def evaluate(model, test_loader):\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs = model(images)\n",
    "            predictions = outputs.argmax(dim=1)\n",
    "            correct += (predictions == labels).sum().item()\n",
    "            total += labels.size(0)\n",
    "    accuracy = correct / total\n",
    "    return accuracy\n",
    "\n",
    "teacher_acc = evaluate(teacher_model, test_loader)\n",
    "print(f\"Teacher Model Accuracy: {teacher_acc:.4f}\")\n",
    "\n",
    "student_acc = evaluate(student_model, test_loader)\n",
    "print(f\"Student Model Accuracy: {student_acc:.4f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
